//
// Created by liugang on 9/12/25.
//
#include "nodes.h"
#include "vars.h"

#include <cmath>

auto spaces(int32_t level) -> const char* {
    static std::string spaces(16, ' ');
    if (spaces.size() < (level * 4)) {
        spaces.resize(level * 4, ' ');
    }

    return (spaces.data() + spaces.size()) - (level * 4);
}

extern auto print_node(const Node* node, int level, uint16_t flags, FILE* file) -> void {
    fflush(file);

    if (node->is_stmt()) {
        char addr_buff[64]{0};
        snprintf(addr_buff, sizeof(addr_buff) - 1, "%p [%c] ", node, node->is_differentiable() ? 'd' : '-');
        fprintf(file, "%-12s", addr_buff);
    }
    fflush(file);

    if (level > 0) {
        fprintf(file, "%s", spaces(level));
        fflush(file);
    }

    switch (node->kind) {
        case NK_EXPR_LITERAL: {
            auto literal = static_cast<const ExprLiteral*>(node);
            fprintf(file, "%lg", literal->val);
        }
        break;
        case NK_EXPR_VAR: {
            auto var = static_cast<const ExprVar*>(node);
            std::string name = var->flow->var_name(var->loc_id);
            fprintf(file, R"(var("%s"))", name.c_str());
        }
        break;
        case NK_EXPR_REF: {
            auto ref = static_cast<const ExprRef*>(node);
            fprintf(file, "ref(%p)", ref->addr);
        }
        break;
        case NK_EXPR_ADD: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " + ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_SUB: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " - ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_MULT: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " * ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_DIV: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " / ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_SIN: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "sin(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_COS: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "cos(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_GT: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " > ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_LT: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " < ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_EQ: {
            auto expr = static_cast<const Expr*>(node);
            fprintf(file, "(");
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " == ");
            print_node(expr->child(1), 0, flags, file);
            fprintf(file, ")");
        }
        break;
        case NK_EXPR_PHI: {
        }
        break;
        case NK_EXPR_ASSIGN: {
            auto expr = static_cast<const Expr*>(node);
            print_node(expr->child(0), 0, flags, file);
            fprintf(file, " = ");
            print_node(expr->child(1), 0, flags, file);
        }
        break;
        case NK_STMT_BLOCK: {
            fprintf(file, "BLOCK");
            fflush(file);
            switch (node->flags & NF_BLOCK_MASK) {
                case NF_BLOCK_IF:
                    fprintf(file, " [if]");
                    break;
                case NF_BLOCK_FOR:
                    fprintf(file, " [for]");
                    break;
                case NF_BLOCK_SEQ:
                    break;
                default:
                    ASSERT(false);
                    break;
            }
            fprintf(file, "\n");

            if (flags & PF_PRINTCHILDREN) {
                LIST_FOREACH(elem, &(node->children)) {
                    print_node(reinterpret_cast<const Node*>(elem), level + 1, flags, file);
                }
            }
        }
        break;
        case NK_STMT_JUMP: {
            const StmtJump* jump = static_cast<const StmtJump*>(node);
            fprintf(file, "JUMP:%p", jump->target);
            switch (node->flags & NF_JUMP_MASK) {
                case NF_JUMP_FORCE: {
                }
                break;
                case NF_JUMP_IF_COND: {
                }
                break;
                case NF_JUMP_FOR_COND: {
                }
                break;
                case NF_JUMP_FOR_BREAK: {
                    fprintf(file, " [break]");
                }
                break;
                case NF_JUMP_FOR_CONTINUE: {
                    fprintf(file, " [continue]");
                }
                break;
                default: {
                    ASSERT(false);
                }
                break;
            }

            // 如果是条件跳转
            if (!list_empty(&jump->children)) {
                fprintf(file, ": ");
                print_node(jump->child(0), 0, flags, file);
            }
            fprintf(file, "\n");
        }
        break;
        case NK_STMT_EXPR: {
            fprintf(file, "(");
            print_node(node->child(0), 0, flags, file);
            fprintf(file, ")");
            fprintf(file, "\n");
        }
        break;
        case NK_STMT_END: {
            fprintf(file, "END\n");
        }
        break;
        case NK_STMT_STOP: {
            fprintf(file, "STOP\n");
        }
        break;
        default: {
            ASSERT(false);
        }
        break;
    }
    fflush(file);
}

extern auto eval_expr(FLOW* flow, Node* n) -> real {
    ASSERT(n->is_expr());

    switch (n->kind) {
        case NK_EXPR_LITERAL: {
            ExprLiteral* node = static_cast<ExprLiteral*>(n);
            node->cached_value = node->val;
        }
        break;
        case NK_EXPR_VAR: {
            ExprVar* var = static_cast<ExprVar*>(n);

            // var->cached_value = var->flow->var_value(var->loc_id);

            FLOW* flow = var->flow;
            AccessProc proc = flow->access(var->access_proc_id);
            n->cached_value = proc(flow->context(), var->loc_id);
        }
        break;
        case NK_EXPR_REF: {
            ExprRef* ref = static_cast<ExprRef*>(n);
            ref->cached_value = *ref->addr;
        }
        break;
        // case NK_EXPR_VAL: {
        //     ExprVal* val = static_cast<ExprVal*>(n);
        //     val->cached_value = *val->addr;
        // }
        // break;
        case NK_EXPR_ADD: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = left_expr->cached_value + right_expr->cached_value;
        }
        break;
        case NK_EXPR_SUB: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = left_expr->cached_value - right_expr->cached_value;
        }
        break;
        case NK_EXPR_MULT: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = left_expr->cached_value * right_expr->cached_value;
        }
        break;
        case NK_EXPR_DIV: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = left_expr->cached_value / right_expr->cached_value;
        }
        break;
        case NK_EXPR_SIN: {
            Node* arg = n->child(0);
            eval_expr(flow, arg);
            n->cached_value = std::sin(arg->cached_value);
        }
        break;
        case NK_EXPR_COS: {
            Node* arg = n->child(0);
            eval_expr(flow, arg);
            n->cached_value = std::cos(arg->cached_value);
        }
        break;
        case NK_EXPR_GT: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = real(left_expr->cached_value > right_expr->cached_value);
        }
        break;
        case NK_EXPR_LT: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = real(left_expr->cached_value < right_expr->cached_value);
        }
        break;
        case NK_EXPR_EQ: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            eval_expr(flow, left_expr);
            eval_expr(flow, right_expr);
            n->cached_value = real(left_expr->cached_value == right_expr->cached_value);
        }
        break;
        case NK_EXPR_PHI: {
            ASSERT(false);
        }
        break;
        case NK_EXPR_ASSIGN: {
            ExprRef* left_expr = static_cast<ExprRef*>(n->child(0));
            Node* right_expr = n->child(1);
            *(left_expr->addr) = eval_expr(flow, right_expr);
            n->cached_value = *(left_expr->addr);
            //fprintf(stdout, "ref(%p) = %lf\n", left_expr->addr, n->cached_value);
        }
        break;
        default: {
            ASSERT(false);
        }
        break;
    }

    auto debug_probe = flow->debug_probe();
    if ((debug_probe != nullptr) && (debug_probe->eval_expr_hook != nullptr)) {
        debug_probe->eval_expr_hook(debug_probe->debug_context, n, n->cached_value);
    }

    return n->cached_value;
}

extern auto eval_stmt(FLOW* flow, Node* root) -> void {
    auto debug_probe = flow->debug_probe();

    // 当前指令
    Node* curr = root;
    while (curr != nullptr) {

        if ((debug_probe != nullptr) && (debug_probe->eval_stmt_hook != nullptr)) {
            debug_probe->eval_stmt_hook(debug_probe->debug_context, curr);
        }

        if (flow)

            switch (curr->kind) {
                case NK_STMT_STOP: {
                    // 停机
                    return;
                }
                case NK_STMT_END: {
                    // END 指令用于表示回到当前指令所述的BLOCK块，
                    // 并取该BLOCK的下一个指令
                    curr = curr->parent;
                    curr = curr->next;
                }
                break;
                case NK_STMT_JUMP: {
                    // JUMP 表示，如果条件表达式的值非0（true），执行下一条指令；
                    // 如果无表达式，或者条件表达式的值为0，跳转到指定的指令；
                    StmtJump* jump = static_cast<StmtJump*>(curr);
                    if (!list_empty(&jump->children)) {
                        // 将跳转条件表达式的值赋值给jump节点
                        jump->cached_value = eval_expr(flow, jump->child(0));
                        fprintf(stdout, "%% = %s\n", ((jump->cached_value != 0) ? "true" : "false"));
                        if (jump->cached_value != 0) {
                            curr = curr->next; // if true, goto next
                            break;
                        }
                    }
                    ASSERT(static_cast<integer>(jump->cached_value) == 0);
                    curr = jump->target;
                }
                break;
                case NK_STMT_BLOCK: {
                    // 表示需要进入BLOCK的子语句，取其第一条一句开始执行
                    curr = curr->child(0);
                }
                break;
                case NK_STMT_EXPR: {
                    eval_expr(flow, curr->child(0));
                    fflush(stdout);
                    curr = curr->next;
                }
                break;
                default: {
                    ASSERT(false);
                }
                break;
            }
    }
}


extern auto deriv_expr(FLOW* flow, Node* n, real derival) -> void {
    ASSERT(n->is_expr());

    // 如果由外层控制语句（比如,for）限制了不能求导就算了
    if (n->is_nondifferentiable()) {
        return;
    }

    switch (n->kind) {
        case NK_EXPR_LITERAL: {
            ASSERT(n->is_nondifferentiable());
        }
        break;
        case NK_EXPR_VAR: {
            auto var = static_cast<ExprVar*>(n);
            flow->var_value(var->loc_id) += derival;
        }
        break;
        case NK_EXPR_REF: {
            ASSERT(n->is_nondifferentiable());
        }
        break;
        case NK_EXPR_ADD: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            deriv_expr(flow, left_expr, derival * 1);
            deriv_expr(flow, right_expr, derival * 1);
        }
        break;
        case NK_EXPR_SUB: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            deriv_expr(flow, left_expr, derival * 1);
            deriv_expr(flow, right_expr, derival * -1);
        }
        break;
        case NK_EXPR_MULT: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            deriv_expr(flow, left_expr, derival * right_expr->cached_value);
            deriv_expr(flow, right_expr, derival * left_expr->cached_value);
        }
        break;
        case NK_EXPR_DIV: {
            Node* left_expr = n->child(0);
            Node* right_expr = n->child(1);
            deriv_expr(flow, left_expr, derival * (1.0 / right_expr->cached_value));
            deriv_expr(flow, right_expr, derival * (-left_expr->cached_value / (left_expr->cached_value * right_expr->cached_value)));
        }
        break;
        case NK_EXPR_SIN: {
            Node* arg = n->child(0);
            deriv_expr(flow, arg, derival * std::cos(arg->cached_value));
        }
        break;
        case NK_EXPR_COS: {
            Node* arg = n->child(0);
            deriv_expr(flow, arg, derival * std::sin(arg->cached_value));
        }
        break;
        case NK_EXPR_GT: {
            ASSERT(n->is_nondifferentiable());
        }
        break;
        case NK_EXPR_LT: {
            ASSERT(n->is_nondifferentiable());
        }
        break;
        case NK_EXPR_EQ: {
            ASSERT(n->is_nondifferentiable());
        }
        break;
        case NK_EXPR_PHI: {
            ASSERT(false);
        }
        break;
        case NK_EXPR_ASSIGN: {
            auto left_expr = static_cast<ExprRef*>(n->child(0));
            Node* right_expr = n->child(1);
            deriv_expr(flow, right_expr, derival);
        }
        break;
        default: {
            ASSERT(false);
        }
        break;
    }
}

extern auto deriv_stmt(FLOW* flow, Node* root, real derval) -> void {
    // 当前指令
    Node* curr = root;
    while (curr != nullptr) {

        auto debug_probe = flow->debug_probe();
        if ((debug_probe != nullptr) && (debug_probe->eval_expr_hook != nullptr)) {
            debug_probe->eval_stmt_hook(debug_probe->debug_context, curr);
        }

        switch (curr->kind) {
            case NK_STMT_STOP: {
                // 停机
                return;
            }
            case NK_STMT_END: {
                // END 指令用于表示回到当前指令所述的BLOCK块，
                // 并取该BLOCK的下一个指令
                curr = curr->parent;
                curr = curr->next;
            }
            break;
            case NK_STMT_JUMP: {
                // JUMP 表示，如果条件表达式的值非0（true），执行下一条指令；
                // 如果无表达式，或者条件表达式的值为0，跳转到指定的指令；
                StmtJump* jump = static_cast<StmtJump*>(curr);
                if (!list_empty(&jump->children)) {
                    if (jump->cached_value != 0) {
                        curr = curr->next; // if true, goto next
                        break;
                    }
                }
                ASSERT(static_cast<integer>(jump->cached_value) == 0);
                curr = jump->target;
            }
            break;
            case NK_STMT_BLOCK: {
                // 表示需要进入BLOCK的子语句，取其第一条一句开始执行
                curr = curr->child(0);
            }
            break;
            case NK_STMT_EXPR: {
                deriv_expr(flow, curr->child(0), derval);
                curr = curr->next;
            }
            break;
            default: {
                ASSERT(false);
            }
            break;
        }
    }
}

